/*
 *  Copyright (C) 2010 Martin Haulrich <mwh.isv@cbs.dk>
 *
 *  This file is part of the MatrixParser package.
 *
 *  The MatrixParser program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Lesser General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package org.osdtsystem.matrixparser.parsers;

import java.util.List;
import org.osdtsystem.matrixparser.data.CONLLSentence;
import org.osdtsystem.matrixparser.featureextraction.FirstOrderExtractorLabelled;
import org.osdtsystem.matrixparser.featureextraction.FirstOrderExtractorUnlabelled;
import org.osdtsystem.matrixparser.featureextraction.UnionExtractor;
import org.osdtsystem.matrixparser.features.FeatureVector;
import org.osdtsystem.matrixparser.features.MstCache;
import org.osdtsystem.matrixparser.features.TwoWayFeatureHandler;
import org.osdtsystem.matrixparser.learners.Scorer;
import org.osdtsystem.matrixparser.features.WeightVector;
import org.osdtsystem.matrixparser.mst.ChuLiuEdmonds;
import org.osdtsystem.matrixparser.parsers.output.DependencyTree;

/**
 * 
 * @author Martin Haulrich
 */
public class LabelledMSTParser extends MSTParser {
    public LabelledMSTParser(String name, Scorer scorer, UnionExtractor extractor, ParsingModel model) {
        this(name, scorer, extractor, extractor, model.labelHandler());
    }

    public LabelledMSTParser(String name, Scorer scorer, FirstOrderExtractorUnlabelled
            unlabelledMstExtractor, FirstOrderExtractorLabelled labelledMstExtractor,
            TwoWayFeatureHandler labelAlphabet) {
        super(name, scorer, unlabelledMstExtractor, labelledMstExtractor, labelAlphabet);
    }

    @Override
    public DependencyTree parse(CONLLSentence sentence, WeightVector weightVector) {
        // Retrieve cache for sentence
        MstCache cache = mstCache(sentence);
        //System.err.println("Weights: " + weightVector.toString());

        // Calculate scores
        int nodes = sentence.size();
        double[][] scores = new double[nodes][nodes];
        int[][] labels = new int[nodes][nodes];
        for (int head = 0; head < nodes; head++) {
            //System.err.print("Best label for " + head + ": ");
            for (int dependent = 0; dependent < nodes; dependent++) {
                double unlabScore = scorer.getScore(weightVector,
                        cache.unlabelledFeatures(head, dependent));

                // Find label with best score;
                double bestScore = Double.MIN_VALUE;
                String bestCheckHead = "";
                String bestCheckDep = "";
                for (int t = 1; t < labelAlphabet.alphabetSize(); t++) {
                    FeatureVector headFV = cache.labelledHeadFeatures(head, dependent, t);
                    FeatureVector depFV = cache.labelledDependentFeatures(head, dependent, t);
                    double score = scorer.getScore(weightVector, headFV) + scorer.getScore(weightVector, depFV);

                    if (score > bestScore) {
                        labels[head][dependent] = t;
                        bestScore = score;
                        bestCheckHead = headFV.checksum64();
                        bestCheckDep = depFV.checksum64();
                    }
                }

                //                    System.err.println(labelAlphabet.getName(types[dependent]) + " " + bestScore);
                //System.err.print("" + dependent + "=" + labels[head][dependent] + "/" + String.format("%.2g", bestScore) + "(" + bestCheckHead + "/" + bestCheckDep + ") ");
                scores[head][dependent] = unlabScore + bestScore;

            }
            //System.err.println();
        }

        // Compute minimum spanning tree
        List<Integer> CLEheads = ChuLiuEdmonds.chuLiuEdmonds(scores);

        // Create solution
        DependencyTree tree = new DependencyTree(sentence, labelAlphabet, lastComplementId);
        for (int node = 0; node < CLEheads.size(); node++) {
            int head = CLEheads.get(node);
            if (head < 0) {
                tree.setDependency(node, head, rootLabel);
            } else {
                int label = labels[head][node];
                tree.setDependency(node, head, label);
            }
        }

        // Return solution
        return tree;
    }

    @Override
    List<FeatureVector> featureVectors(DependencyTree tree, List<FeatureVector> aggregator) {
        super.featureVectors(tree, aggregator);
        CONLLSentence sentence = tree.sentence();
        MstCache cache = mstCache(sentence);
        for (int node = 1; node < tree.nodes(); ++ node) {
            int head = tree.head(node);
            int label = tree.label(node);
            if (head >= 0) {
                aggregator.add(cache.labelledDependentFeatures(head, node, label));
                aggregator.add(cache.labelledHeadFeatures(head, node, label));
            }
        }
        return aggregator;
    }
}
